import re
import json
from Bio.SeqUtils.ProtParam import ProteinAnalysis

def write_json(file_name,data):
    with open(file_name, 'w') as json_file:
        json.dump(data, json_file)

def read(filename): # 读取数据
    seqs=list()
    with open(filename,"r",encoding="utf8",) as f:
        for line in f :
            line=line.strip()
            if line[0] != ">":  # 选择序列行
                seqs.append(line)
    return seqs

def cleave(seqs):   # 正则比配剪切
    peptides=dict()
    trypsin_pattern = re.compile(r'(?<=[RK])(?!P)')
    for seq in seqs:
        peptides[seq]=list()
        for peptide in trypsin_pattern.split(seq):
            peptides[seq].append({"sequences":peptide})
    return peptides

def get_weight(seq):
    weight=ProteinAnalysis(seq).molecular_weight()
    weight+=seq.count('C')*57.021
    return weight

def cleave_ex1(filename):
    seqs=read(filename=filename)
    return cleave(seqs)

def fragmentation(seq):
    result={}
    for Q in (2,3):
        result[f"with {Q} charge"]=[]
        if len(seq) == 1:
            weight=get_weight(seq)+Q
            result[f"with {Q} charge"].append({
                "type":"None",
                "sequences":seq,
                "weight":weight,
                "chrage":Q,
                "m/z":weight/Q
            })
            break
        for i in range(len(seq)-1):
            if Q == 2:
                Qs=[(1,1)]
            if Q == 3:
                Qs=[(1,2),(2,1)]
            for q1,q2 in Qs:
                weight=get_weight(seq[:i+1])+q1-18
                result[f"with {Q} charge"].append({
                    "type":f"b{i+1}",
                    "sequences":seq[:i+1],
                    "weight":weight,
                    "chrage":q1,
                    "m/z":weight/q1
                })
                weight=get_weight(seq[i+1:])+q2
                result[f"with {Q} charge"].append({
                    "type":f"y{len(seq)-i-1}",
                    "sequences":seq[i+1:],
                    "weight":weight,
                    "chrage":q2,
                    "m/z":weight/q2
                })
    return result

if __name__ == "__main__":
    seqs_dict=cleave_ex1("P01308.fasta")
    write_json("test.json",seqs_dict)
    for peptide,cleave_parts in seqs_dict.items():
        for index,cleave_part in enumerate(cleave_parts):
            chraged=fragmentation(cleave_part["sequences"])
            seqs_dict[peptide][index].update(chraged)
    write_json("理论by离子的比荷.json",seqs_dict)
